{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Intro to Batch Predictions with the Gemini API\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fbatch-prediction%2Fintro_batch_prediction.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/batch-prediction/intro_batch_prediction.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Eric Dong](https://github.com/gericdong), [Holt Skinner](https://github.com/holtskinner) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Different from getting online (synchronous) responses, where you are limited to one input request at a time, the batch predictions with the Gemini API in Vertex AI allow you to send a large number of multimodal requests to a Gemini model in a single batch request. Then, the model responses asynchronously populate to your storage output location in [Cloud Storage](https://cloud.google.com/storage/docs/introduction) or [BigQuery](https://cloud.google.com/bigquery/docs/storage_overview).\n",
        "\n",
        "Batch predictions are generally more efficient and cost-effective than online predictions when processing a large number of inputs that are not latency sensitive.\n",
        "\n",
        "To learn more, see the [Get batch predictions for Gemini](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini) page.\n",
        "\n",
        "### Objectives\n",
        "\n",
        "In this tutorial, you learn how to make batch predictions with the Gemini API in Vertex AI. This tutorial shows how to use **Cloud Storage** and **BigQuery** as input sources and output locations.\n",
        "\n",
        "You will complete the following tasks:\n",
        "\n",
        "- Preparing batch inputs and an output location\n",
        "- Submitting a batch prediction job\n",
        "- Retrieving batch prediction results\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai pandas google-cloud-storage google-cloud-bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06489bd14f16"
      },
      "source": [
        "### Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fe00fa0b8bb7"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "import time\n",
        "\n",
        "import fsspec\n",
        "from google import genai\n",
        "from google.cloud import bigquery\n",
        "from google.genai.types import CreateBatchJobConfig\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and create client\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"global\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cfca4d7bd6db"
      },
      "outputs": [],
      "source": [
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e43229f3ad4f"
      },
      "source": [
        "### Load model\n",
        "\n",
        "You can find a list of the Gemini models that support batch predictions in the [Multimodal models that support batch predictions](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini#multimodal_models_that_support_batch_predictions) page.\n",
        "\n",
        "This tutorial uses Gemini 2.5 Flash (`gemini-2.5-flash`) model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf93d5f0ce00"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.5-flash\"  # @param {type:\"string\", isTemplate: true}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "265f180b58e0"
      },
      "source": [
        "## Cloud Storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1_xZADsak23H"
      },
      "source": [
        "### Prepare batch inputs\n",
        "\n",
        "The input for batch requests specifies the items to send to your model for prediction. You can learn more about the batch input formats in the [Batch text generation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini#prepare_your_inputs) page.\n",
        "\n",
        "This tutorial uses Cloud Storage as an example. The requirements for Cloud Storage input are:\n",
        "\n",
        "- File format: [JSON Lines (JSONL)](https://jsonlines.org/)\n",
        "- Located in `us-central1`\n",
        "- Appropriate read permissions for the service account\n",
        "\n",
        "Each request that you send to a model can include parameters that control how the model generates a response. Learn more about Gemini parameters in the [Experiment with parameter values](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/adjust-parameter-values) page.\n",
        "\n",
        "This is one of the example requests in the input JSONL file `batch_requests_for_multimodal_input_2.jsonl`:\n",
        "\n",
        "```json\n",
        "{\"request\":{\"contents\": [{\"role\": \"user\", \"parts\": [{\"text\": \"List objects in this image.\"}, {\"file_data\": {\"file_uri\": \"gs://cloud-samples-data/generative-ai/image/office-desk.jpeg\", \"mime_type\": \"image/jpeg\"}}]}],\"generationConfig\":{\"temperature\": 0.4}}}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uWb8QzxwbH6W"
      },
      "outputs": [],
      "source": [
        "INPUT_DATA = \"gs://cloud-samples-data/generative-ai/batch/batch_requests_for_multimodal_input_2.jsonl\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T3jQ59mCsXLc"
      },
      "source": [
        "### Prepare batch output location\n",
        "\n",
        "When a batch prediction task completes, the output is stored in the location that you specified in your request.\n",
        "\n",
        "- The location is in the form of a Cloud Storage prefix.\n",
        "  - For example: `gs://path/to/output/data`.\n",
        "\n",
        "- You can specify the URI of your Cloud Storage bucket in `BUCKET_URI`, or\n",
        "- If it is not specified, this notebook will create a Cloud Storage bucket in the form of `gs://PROJECT_ID-TIMESTAMP`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OtUodwGXZ7US"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = \"[your-cloud-storage-bucket]\"  # @param {type:\"string\"}\n",
        "GCS_LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "if BUCKET_URI == \"[your-cloud-storage-bucket]\":\n",
        "    TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-{TIMESTAMP}\"\n",
        "\n",
        "    ! gsutil mb -l {GCS_LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T90CwWDHvonn"
      },
      "source": [
        "### Send a batch prediction request\n",
        "\n",
        "To make a batch prediction request, you specify a source model ID, an input source and an output location where Vertex AI stores the batch prediction results.\n",
        "\n",
        "To learn more, see the [Batch prediction API](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/batch-prediction-api) page.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d8e54c57072e"
      },
      "outputs": [],
      "source": [
        "gcs_batch_job = client.batches.create(\n",
        "    model=MODEL_ID,\n",
        "    src=INPUT_DATA,\n",
        "    config=CreateBatchJobConfig(dest=BUCKET_URI),\n",
        ")\n",
        "gcs_batch_job.name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-Fo_Kd9FYRj"
      },
      "source": [
        "Print out the job status and other properties. You can also check the status in the Cloud Console at https://console.cloud.google.com/vertex-ai/batch-predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DWq7m79PbjG8"
      },
      "outputs": [],
      "source": [
        "gcs_batch_job = client.batches.get(name=gcs_batch_job.name)\n",
        "gcs_batch_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WHUUEREoiewD"
      },
      "source": [
        "Optionally, you can list all the batch prediction jobs in the project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QVgOnasfigx1"
      },
      "outputs": [],
      "source": [
        "for job in client.batches.list():\n",
        "    print(job.name, job.create_time, job.state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aJaPNBrGPqK"
      },
      "source": [
        "### Wait for the batch prediction job to complete\n",
        "\n",
        "Depending on the number of input items that you submitted, a batch generation task can take some time to complete. You can use the following code to check the job status and wait for the job to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dtJDIXdHc0W-"
      },
      "outputs": [],
      "source": [
        "# Refresh the job until complete\n",
        "while gcs_batch_job.state in (\n",
        "    \"JOB_STATE_RUNNING\",\n",
        "    \"JOB_STATE_PENDING\",\n",
        "    \"JOB_STATE_QUEUED\",\n",
        "):\n",
        "    time.sleep(5)\n",
        "    gcs_batch_job = client.batches.get(name=gcs_batch_job.name)\n",
        "\n",
        "# Check if the job succeeds\n",
        "if gcs_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
        "    print(\"Job succeeded!\")\n",
        "else:\n",
        "    print(f\"Job failed: {gcs_batch_job.error}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XWUgAxL-HjN9"
      },
      "source": [
        "### Retrieve batch prediction results\n",
        "\n",
        "When a batch prediction task is complete, the output of the prediction is stored in the bucket in JSONL that you specified in your request.\n",
        "\n",
        "The file name should look like this: `{gcs_batch_job.dest.gcs_uri}/prediction-model-TIMESTAMP/predictions.jsonl`\n",
        "\n",
        "Example output:\n",
        "\n",
        "```json\n",
        "{\"status\": \"\", \"processed_time\": \"2024-11-13T14:04:28.376+00:00\", \"request\": {\"contents\": [{\"parts\": [{\"file_data\": null, \"text\": \"List objects in this image.\"}, {\"file_data\": {\"file_uri\": \"gs://cloud-samples-data/generative-ai/image/gardening-tools.jpeg\", \"mime_type\": \"image/jpeg\"}, \"text\": null}], \"role\": \"user\"}], \"generationConfig\": {\"temperature\": 0.4}}, \"response\": {\"candidates\": [{\"avgLogprobs\": -0.10394711927934126, \"content\": {\"parts\": [{\"text\": \"Here's a list of the objects in the image:\\n\\n* **Watering can:** A green plastic watering can with a white rose head.\\n* **Plant:** A small plant (possibly oregano) in a terracotta pot.\\n* **Terracotta pots:** Two terracotta pots, one containing the plant and another empty, stacked on top of each other.\\n* **Gardening gloves:** A pair of striped gardening gloves.\\n* **Gardening tools:** A small trowel and a hand cultivator (hoe).  Both are green with black handles.\"}], \"role\": \"model\"}, \"finishReason\": \"STOP\"}], \"modelVersion\": \"gemini-2.0-flash@default\", \"usageMetadata\": {\"candidatesTokenCount\": 110, \"promptTokenCount\": 264, \"totalTokenCount\": 374}}}\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qZlRIYsC01F1"
      },
      "source": [
        "The example code below shows how to load the `.jsonl` file in the Cloud Storage output location into a Pandas DataFrame and print out the object.\n",
        "\n",
        "You can retrieve the specific responses in the `response` field."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-jLl3es3dTqB"
      },
      "outputs": [],
      "source": [
        "fs = fsspec.filesystem(\"gcs\")\n",
        "\n",
        "file_paths = fs.glob(f\"{gcs_batch_job.dest.gcs_uri}/*/predictions.jsonl\")\n",
        "\n",
        "if gcs_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
        "    # Load the JSONL file into a DataFrame\n",
        "    df = pd.read_json(f\"gs://{file_paths[0]}\", lines=True)\n",
        "\n",
        "    df = df.join(pd.json_normalize(df[\"response\"], \"candidates\"))\n",
        "    display(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfb2a462a7c6"
      },
      "source": [
        "## BigQuery\n",
        "\n",
        "⚠️ Batch predictions using BigQuery currently does not support Global endpoints. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c694724e069"
      },
      "outputs": [],
      "source": [
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=\"us-central1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ea69e283023"
      },
      "source": [
        "### Batch Input Preparation  \n",
        "\n",
        "To send batch requests for prediction, you need to structure your input properly. For more details, visit the [Batch text generation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini#prepare_your_inputs) page.  \n",
        "\n",
        "This guide uses **BigQuery** as an example. To use a BigQuery table as input:  \n",
        "- Ensure the dataset is created in a supported region (e.g., `us-central1`). Multi-region locations (e.g., `us`) are not allowed.  \n",
        "- The input table must include a `request` column of type `JSON` or `STRING` containing valid JSON, structured as a [`GenerateContentRequest`](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference).  \n",
        "- Additional columns can use any BigQuery data types except `array`, `struct`, `range`, `datetime`, and `geography`. These are ignored for generation but appear in the output table. The system reserves `response` and `status` for output.  \n",
        "- Only public YouTube or Cloud Storage URIs are supported in the `fileData` or `file_data` field.  \n",
        "- Requests can include parameters to customize the model's output. Learn more in the [Gemini parameters guide](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/adjust-parameter-values)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "835bfd485d49"
      },
      "source": [
        "This is an example BigQuery table with sample requests:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b932f02d1f9c"
      },
      "outputs": [],
      "source": [
        "INPUT_DATA = \"bq://storage-samples.generative_ai.batch_requests_for_multimodal_input_2\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d62eb771ad6"
      },
      "source": [
        "You can query the BigQuery table to review the input data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fbfcefb7d295"
      },
      "outputs": [],
      "source": [
        "bq_client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "bq_table_id = INPUT_DATA.replace(\"bq://\", \"\")\n",
        "sql = f\"\"\"\n",
        "        SELECT *\n",
        "        FROM {bq_table_id}\n",
        "        \"\"\"\n",
        "\n",
        "query_result = bq_client.query(sql)\n",
        "\n",
        "df = query_result.result().to_dataframe()\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "529dae543bfc"
      },
      "source": [
        "### Prepare batch output location\n",
        "\n",
        "When a batch prediction task completes, the output is stored in the location that you specified in your request.\n",
        "\n",
        "- The location is in the form of a BigQuery URI prefix, for example: `bq://projectId.bqDatasetId`.\n",
        "- If not specified, `bq://PROJECT_ID.gen_ai_batch_prediction.predictions_TIMESTAMP` will be used.\n",
        "\n",
        "This tutorial uses a **BigQuery** table as an example.\n",
        "\n",
        "- You can specify the URI of your BigQuery table in `BQ_OUTPUT_URI`, or\n",
        "- If it is not specified, this notebook will create a new dataset `bq://PROJECT_ID.gen_ai_batch_prediction` for you."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "75914840e555"
      },
      "outputs": [],
      "source": [
        "BQ_OUTPUT_URI = \"[your-bigquery-table]\"  # @param {type:\"string\"}\n",
        "\n",
        "if BQ_OUTPUT_URI == \"[your-bigquery-table]\":\n",
        "    bq_dataset_id = \"gen_ai_batch_prediction\"\n",
        "\n",
        "    # The output table will be created automatically if it doesn't exist\n",
        "    timestamp = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "    bq_table_id = f\"prediction_result_{timestamp}\"\n",
        "    BQ_OUTPUT_URI = f\"bq://{PROJECT_ID}.{bq_dataset_id}.{bq_table_id}\"\n",
        "\n",
        "    bq_dataset = bigquery.Dataset(f\"{PROJECT_ID}.{bq_dataset_id}\")\n",
        "    bq_dataset.location = \"us-central1\"\n",
        "\n",
        "    bq_dataset = bq_client.create_dataset(bq_dataset, exists_ok=True, timeout=30)\n",
        "    print(\n",
        "        f\"Created BigQuery dataset {bq_client.project}.{bq_dataset.dataset_id} for batch prediction output.\"\n",
        "    )\n",
        "\n",
        "print(f\"BigQuery output URI: {BQ_OUTPUT_URI}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "14f84475dc26"
      },
      "source": [
        "### Send a batch prediction request\n",
        "\n",
        "To make a batch prediction request, you specify a source model ID, an input source and an output location where Vertex AI stores the batch prediction results.\n",
        "\n",
        "To learn more, see the [Batch prediction API](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/batch-prediction-api) page.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e809955e753b"
      },
      "outputs": [],
      "source": [
        "bq_batch_job = client.batches.create(\n",
        "    model=MODEL_ID,\n",
        "    src=INPUT_DATA,\n",
        "    config=CreateBatchJobConfig(dest=BQ_OUTPUT_URI),\n",
        ")\n",
        "bq_batch_job.name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fcb3dc2a5fc"
      },
      "source": [
        "Print out the job status and other properties. You can also check the status in the Cloud Console at https://console.cloud.google.com/vertex-ai/batch-predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b19319d92bf0"
      },
      "outputs": [],
      "source": [
        "bq_batch_job = client.batches.get(name=bq_batch_job.name)\n",
        "bq_batch_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e72800144910"
      },
      "source": [
        "Optionally, you can list all the batch prediction jobs in the project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6a9fb087f9ba"
      },
      "outputs": [],
      "source": [
        "for job in client.batches.list():\n",
        "    print(job.name, job.create_time, job.state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ebcb00b3add9"
      },
      "source": [
        "### Wait for the batch prediction job to complete\n",
        "\n",
        "Depending on the number of input items that you submitted, a batch generation task can take some time to complete. You can use the following code to check the job status and wait for the job to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "189945468c6b"
      },
      "outputs": [],
      "source": [
        "# Refresh the job until complete\n",
        "while bq_batch_job.state in (\n",
        "    \"JOB_STATE_RUNNING\",\n",
        "    \"JOB_STATE_PENDING\",\n",
        "    \"JOB_STATE_QUEUED\",\n",
        "):\n",
        "    time.sleep(5)\n",
        "    bq_batch_job = client.batches.get(name=bq_batch_job.name)\n",
        "\n",
        "# Check if the job succeeds\n",
        "if bq_batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
        "    print(\"Job succeeded!\")\n",
        "else:\n",
        "    print(f\"Job failed: {bq_batch_job.error}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ec25f40f0dd9"
      },
      "source": [
        "### Retrieve batch prediction results\n",
        "\n",
        "When a batch prediction task is complete, the output of the prediction is stored in the location that you specified in your request. It is also available in `batch_job.dest.bigquery_uri` or `batch_job.dest.gcs_uri`.\n",
        "\n",
        "- When you are using BigQuery, the output of batch prediction is stored in an output dataset. If you had provided a dataset, the name of the dataset (`BQ_OUTPUT_URI`) is the name you had provided earlier. \n",
        "- If you did not provide an output dataset, a default dataset `bq://PROJECT_ID.gen_ai_batch_prediction` will be created for you.\n",
        "- The name of the table is formed by appending `predictions_` with the timestamp of when the batch prediction job started."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c459121f0169"
      },
      "source": [
        "You can use the example code below to retrieve predictions and store them into a Pandas DataFrame.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ca8b73b0ad1b"
      },
      "outputs": [],
      "source": [
        "bq_table_id = bq_batch_job.dest.bigquery_uri.replace(\"bq://\", \"\")\n",
        "\n",
        "sql = f\"\"\"\n",
        "        SELECT *\n",
        "        FROM {bq_table_id}\n",
        "        \"\"\"\n",
        "\n",
        "query_result = bq_client.query(sql)\n",
        "\n",
        "df = query_result.result().to_dataframe()\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Clean up resources created in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZNCyIKIrdPJY"
      },
      "outputs": [],
      "source": [
        "# Delete the batch prediction jobs\n",
        "if gcs_batch_job:\n",
        "    client.batches.delete(name=gcs_batch_job.name)\n",
        "if bq_batch_job:\n",
        "    client.batches.delete(name=bq_batch_job.name)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_batch_prediction.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
